{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Private AI PII detection example\n",
    "\n",
    "This notebook shows how to use Private AI for PII detection and PII masking in NeMo Guardrails."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PII Detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from nemoguardrails import LLMRails, RailsConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create rails with Private AI PII detection\n",
    "\n",
    "For this step you'll need your OpenAI API key & Private AI API key.\n",
    "\n",
    "You can get your Private AI API key by signing up on the [Private AI Portal](https://portal.private-ai.com). For more details on Private AI integration, check out this [user guide](../../docs/user-guides/community/privateai.md).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PAI_API_KEY\"] = \"YOUR PRIVATE AI API KEY\"  # Visit https://portal.private-ai.com to get your API key\n",
    "\n",
    "YAML_CONFIG = \"\"\"\n",
    "models:\n",
    "  - type: main\n",
    "    engine: openai\n",
    "    model: gpt-3.5-turbo-instruct\n",
    "\n",
    "rails:\n",
    "  config:\n",
    "    privateai:\n",
    "      server_endpoint: https://api.private-ai.com/cloud/v3/process/text\n",
    "      input:\n",
    "        entities:\n",
    "          - NAME_FAMILY\n",
    "          - LOCATION_ADDRESS_STREET\n",
    "          - EMAIL_ADDRESS\n",
    "      output:\n",
    "        entities:\n",
    "          - NAME_FAMILY\n",
    "          - LOCATION_ADDRESS_STREET\n",
    "          - EMAIL_ADDRESS\n",
    "  input:\n",
    "    flows:\n",
    "      - detect pii on input\n",
    "\n",
    "  output:\n",
    "    flows:\n",
    "      - detect pii on output\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "config = RailsConfig.from_content(yaml_content=YAML_CONFIG)\n",
    "rails = LLMRails(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input rails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = rails.generate(\n",
    "    messages=[{\"role\": \"user\", \"content\": \"Hello! I'm John. My email id is text@gmail.com. I live in California, USA.\"}]\n",
    ")\n",
    "\n",
    "info = rails.explain()\n",
    "\n",
    "print(\"Response\")\n",
    "print(\"----------------------------------------\")\n",
    "print(response[\"content\"])\n",
    "\n",
    "\n",
    "print(\"\\n\\nColang history\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.colang_history)\n",
    "\n",
    "print(\"\\n\\nLLM calls summary\")\n",
    "print(\"----------------------------------------\")\n",
    "info.print_llm_calls_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output rails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"give me a sample email id\"}])\n",
    "\n",
    "info = rails.explain()\n",
    "\n",
    "print(\"Response\")\n",
    "print(\"----------------------------------------\\n\\n\")\n",
    "print(response[\"content\"])\n",
    "\n",
    "\n",
    "print(\"\\n\\nColang history\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.colang_history)\n",
    "\n",
    "print(\"\\n\\nLLM calls summary\")\n",
    "print(\"----------------------------------------\")\n",
    "info.print_llm_calls_summary()\n",
    "\n",
    "\n",
    "print(\"\\n\\nCompletions where PII was detected!\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.llm_calls[0].completion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PII Masking\n",
    "\n",
    "Note: This example uses ollama model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input rails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PAI_API_KEY\"] = \"YOUR PRIVATE AI API KEY\"  # Visit https://portal.private-ai.com to get your API key\n",
    "\n",
    "YAML_CONFIG = \"\"\"\n",
    "<!-- models:\n",
    "  - type: main\n",
    "    engine: openai\n",
    "    model: gpt-3.5-turbo-instruct -->\n",
    "\n",
    "models:\n",
    "  - type: main\n",
    "    engine: ollama\n",
    "    model: llama3.2\n",
    "    parameters:\n",
    "      base_url: http://localhost:11434\n",
    "\n",
    "rails:\n",
    "  config:\n",
    "    privateai:\n",
    "      server_endpoint: https://api.private-ai.com/cloud/v3/process/text\n",
    "      input:\n",
    "        entities:\n",
    "          - LOCATION\n",
    "          - EMAIL_ADDRESS\n",
    "  input:\n",
    "    flows:\n",
    "      - mask pii on input\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "config = RailsConfig.from_content(yaml_content=YAML_CONFIG)\n",
    "rails = LLMRails(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = rails.generate(\n",
    "    messages=[{\"role\": \"user\", \"content\": \"Hello! I'm John. My email id is text@gmail.com. I live in California, USA.\"}]\n",
    ")\n",
    "\n",
    "info = rails.explain()\n",
    "\n",
    "print(\"Response\")\n",
    "print(\"----------------------------------------\")\n",
    "print(response[\"content\"])\n",
    "\n",
    "\n",
    "print(\"\\n\\nColang history\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.colang_history)\n",
    "\n",
    "print(\"\\n\\nLLM calls summary\")\n",
    "print(\"----------------------------------------\")\n",
    "info.print_llm_calls_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output rails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PAI_API_KEY\"] = \"YOUR PRIVATE AI API KEY\"  # Visit https://portal.private-ai.com to get your API key\n",
    "\n",
    "YAML_CONFIG = \"\"\"\n",
    "<!-- models:\n",
    "  - type: main\n",
    "    engine: openai\n",
    "    model: gpt-3.5-turbo-instruct -->\n",
    "\n",
    "models:\n",
    "  - type: main\n",
    "    engine: ollama\n",
    "    model: llama3.2\n",
    "    parameters:\n",
    "      base_url: http://localhost:11434\n",
    "\n",
    "rails:\n",
    "  config:\n",
    "    privateai:\n",
    "      server_endpoint: https://api.private-ai.com/cloud/v3/process/text\n",
    "      output:\n",
    "        entities:\n",
    "          - LOCATION\n",
    "          - EMAIL_ADDRESS\n",
    "  output:\n",
    "    flows:\n",
    "      - mask pii on output\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "config = RailsConfig.from_content(yaml_content=YAML_CONFIG)\n",
    "rails = LLMRails(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"give me a sample email id\"}])\n",
    "\n",
    "info = rails.explain()\n",
    "\n",
    "print(\"Response\")\n",
    "print(\"----------------------------------------\\n\\n\")\n",
    "print(response[\"content\"])\n",
    "\n",
    "\n",
    "print(\"\\n\\nColang history\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.colang_history)\n",
    "\n",
    "print(\"\\n\\nLLM calls summary\")\n",
    "print(\"----------------------------------------\")\n",
    "info.print_llm_calls_summary()\n",
    "\n",
    "\n",
    "print(\"\\n\\nCompletions where PII was detected!\")\n",
    "print(\"----------------------------------------\")\n",
    "print(info.llm_calls[0].completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nemo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
